#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
All rights reserved.
2010.07.29
''' 

from msexcel import *
from textcolor import DCC
from main import Config, info_from_xlsx
import main
import sys
import os
import os.path

def print_error(str):
    DCC.setColor(DCC.red)
    print str
    DCC.resetColor()

def print_info(str, s2=""):
    DCC.setColor(DCC.green)
    if s2 == ',':
        print str,
    elif s2:
        print str, s2
    else:
        print str
    DCC.resetColor()

def readResultFiles(wd_list, le_list, ap_list):
    data = []
    for wd in wd_list:
        data_wd = []
        for le in le_list:
            data_le = readResultFile(wd, le, len(ap_list))
            data_wd.append(data_le)
        data.append(data_wd)
    return data

def readResultFile(wd, le, ap_list_size):
    u"""
    工作距离 wd，落点能量le
    """
    result_file = main.gen_result_file_name(wd,le)
    if Config.create_sub_folder:
        wd_dir = main.gen_wd_dir_name(wd)
        le_dir = main.gen_le_dir_name(le)
        result_file = os.path.join(wd_dir, le_dir, result_file)
    result_file = os.path.abspath(result_file)
    data = []
    if not os.path.exists(result_file):
        print_error(u"文件 %s 未找到" % result_file)
        print_error(u"继续吗？..")
        if not Config.is_demo:
            raw_input()
        data = [(None, None, None)]*ap_list_size
    else:
        print_info(result_file)
        with open(result_file) as f:
            lines = f.readlines()
        for index in range(ap_list_size):
            try:
                items = lines[index].split()
            except:
                data.append((None, None, None))
                continue
            if len(items) != 4:
                data.append((None, None, None))
                continue
            try:
                v1 = float(items[1])
            except:
                v1 = None
            try:
                v2 = float(items[2])
            except:
                v2 = None
            try:
                v3 = float(items[3])
            except:
                v3 = None
            data.append((v1, v2, v3))
    return data

def getValue(wd, le, ap_index):
    u"""
    获取 工作距离 wd，落点能量le，光阑大小 ap下的值
    一次获取一个值，实现简单，但不是太高效
    """
    pass


def extract_other_two(first, elens_list, le_list, data):
    u"""
    提取激励和灵敏度
    """
    excitation = []
    sensitivities = []

    if first == 0: #工作距离第一行，落点能量为列
        line_items = ["%15s" % " LE"]
        for i in range(len(elens_list)): # 输出各个工作距离
            line_items.append("%15s" % elens_list[i])
        line_items.append('\n')
        excitation.append(' '.join(line_items)) # header line
        sensitivities.append(' '.join(line_items)) # header line
        for row in range(len(le_list)): # 对每个落点能量，输出
            line_items = ["%15s" % le_list[row]]
            for col in range(len(elens_list)):
                line_items.append("%15s" % data[col][row][0][0])
            line_items.append('\n')
            excitation.append(' '.join(line_items))

            line_items = ["%15s" % le_list[row]]
            for col in range(len(elens_list)):
                line_items.append("%15s" % data[col][row][0][1])
            line_items.append('\n')
            sensitivities.append(' '.join(line_items))
    else: #落点能量第一行， 每行 对应不同工作距离
        line_itmes = ["%15s" % "WD"]
        for i in range(len(le_list)):
            line_itmes.append("%15s" % le_list[i])
        line_items.append('\n')
        excitation.append(' '.join(line_items)) # header line
        sensitivities.append(' '.join(line_items)) # header line
        for row in range(len(elens_list)):
            line_items = ["%15s" % elens_list[row]]
            for col in range(len(le_list)):
                line_items.append("%15s" % data[row][col][0][0])
            line_items.append('\n')
            excitation.append(' '.join(line_items))

            line_items = ["%15s" % elens_list[row]]
            for col in range(len(le_list)):
                line_items.append("%15s" % data[row][col][0][1])
            line_items.append('\n')
            sensitivities.append(' '.join(line_items))

    filename = u"激励文件.txt"
    with open(filename, 'w') as f:
        f.writelines(excitation)
        print_info(u"保存：%s" % filename)
    filename = u"偏转灵敏度文件.txt"
    with open(filename, 'w') as f:
        f.writelines(sensitivities)
        print_info(u"保存：%s" % filename)

def dealwith_aperture_xlsx(first, second, elens_list, le_list, aperture_list, data):
    u'''
    按光阑大小生成xlsx文件
    '''
    for ap_index in range(len(aperture_list)):
        xls = MSExcel()
        if first == 0: # 工作距离第一行
            for i in range(len(elens_list)):
                xls.setcell(1, 2+i, elens_list[i], 1)
                xls.setcell(1, 2+i, elens_list[i], 2)
                xls.setcell(1, 2+i, elens_list[i], 3)
            for i in range(len(le_list)):
                xls.setcell(2+i, 1, le_list[i], 1)
                xls.setcell(2+i, 1, le_list[i], 2)
                xls.setcell(2+i, 1, le_list[i], 3)
        else:
            for i in range(len(elens_list)):
                xls.setcell(2+i, 1, elens_list[i], 1)
                xls.setcell(2+i, 1, elens_list[i], 2)
                xls.setcell(2+i, 1, elens_list[i], 3)
            for i in range(len(le_list)):
                xls.setcell(1, 2+i, le_list[i], 1)
                xls.setcell(1, 2+i, le_list[i], 2)
                xls.setcell(1, 2+i, le_list[i], 3)

        for wd_index in range(len(elens_list)):
            for le_index in range(len(le_list)):
                v1, v2, v3 = data[wd_index][le_index][ap_index]
                if first == 0: #工作距离第一行， 行 对应不同落点能量
                    xls.setcell(le_index+2, wd_index+2, v1 if v1 else 0, 1)
                    xls.setcell(le_index+2, wd_index+2, v2 if v2 else 0, 2)
                    xls.setcell(le_index+2, wd_index+2, v3 if v3 else 0, 3)
                    print_info("working... %s %s %s" % (elens_list[wd_index], le_list[le_index], aperture_list[ap_index]))
                else:
                    xls.setcell(wd_index+2, le_index+2, v1 if v1 else 0, 1)
                    xls.setcell(wd_index+2, le_index+2, v2 if v2 else 0, 2)
                    xls.setcell(wd_index+2, le_index+2, v3 if v3 else 0, 3)

        xls.save(u"光阑大小%s.xlsx" % aperture_list[ap_index])
        print_info(u"保存：光阑大小%s.xlsx" % aperture_list[ap_index])
        del xls

def dealwith_aperture(first, second, elens_list, le_list, aperture_list, data):
    u'''
    按光阑大小生成 .txt文件
    '''
    for ap_index in range(len(aperture_list)):
        file_data = []
        if first == 0: #工作距离第一行， 每行 对应不同落点能量
            line_items = ["%15s" % " LE"]
            for i in range(len(elens_list)): # 输出各个工作距离
                line_items.append("%15s" % elens_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(le_list)): # 对每个落点能量，输出
                line_items = ["%15s" % le_list[row]]
                for col in range(len(elens_list)):
                    line_items.append("%15s" % data[col][row][ap_index][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))
        else: #落点能量第一行， 每行 对应不同工作距离
            line_itmes = ["%15s" % "WD"]
            for i in range(len(le_list)):
                line_itmes.append("%15s" % le_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(elens_list)):
                line_items = ["%15s" % elens_list[row]]
                for col in range(len(le_list)):
                    line_items.append("%15s" % data[row][col][ap_index][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))

        filename = u"光阑大小%s.txt" % aperture_list[ap_index]
        with open(filename, 'w') as f:
            f.writelines(file_data)
            print_info(u"保存：%s" % filename)

def dealwith_le(first, second, elens_list, le_list, aperture_list, data):
    u'''
    按落点能量生成 .txt 文件
    '''
    for le_index in range(len(le_list)):
        file_data = []
        if first == 0: # 工作距离第一行， 每行 对应不同光阑大小 
            line_items = ["%15s" % "aperture"]
            for i in range(len(elens_list)):
                line_items.append("%15s" % elens_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(aperture_list)):
                line_items = ["%15s" % aperture_list[row]]
                for col in range(len(elens_list)):
                    line_items.append("%15s" % data[col][le_index][row][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))
        else:
            line_itmes = ["%15s" % "WD"]
            for i in range(len(aperture_list)):
                line_itmes.append("%15s" % aperture_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(elens_list)):
                line_items = ["%15s" % elens_list[row]]
                for col in range(len(aperture_list)):
                    line_items.append("%15s" % data[row][le_index][col][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))

        filename = u"落点能量%s.txt" % le_list[le_index]
        with open(filename, 'w') as f:
            f.writelines(file_data)
            print_info(u"保存：%s" % filename)

def dealwith_wd(first, second, elens_list, le_list, aperture_list, data):
    u'''
    按工作距离生成 .txt 文件
    '''
    for wd_index in range(len(elens_list)):
        file_data = []
        if first == 1: # 落点能量为第一行， 每行 对应不同光阑大小
            line_items = ["%15s" % "aperture"]
            for i in range(len(le_list)):
                line_items.append("%15s" % le_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(aperture_list)):
                line_items = ["%15s" % aperture_list[row]]
                for col in range(len(le_list)):
                    line_items.append("%15s" % data[wd_index][col][row][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))
        else:
            line_itmes = ["%15s" % "WD"]
            for i in range(len(aperture_list)):
                line_itmes.append("%15s" % aperture_list[i])
            line_items.append('\n')
            file_data.append(' '.join(line_items)) # header line
            for row in range(len(le_list)):
                line_items = ["%15s" % le_list[row]]
                for col in range(len(aperture_list)):
                    line_items.append("%15s" % data[wd_index][row][col][2])
                line_items.append('\n')
                file_data.append(' '.join(line_items))

        filename = u"工作距离%s.txt" % elens_list[wd_index]
        with open(filename, 'w') as f:
            f.writelines(file_data)
            print_info(u"保存：%s" % filename)

if __name__ == '__main__':
    vals = [u"工作距离", u"落点能量", u"光阑大小"]
    xlsfile = MSExcel(Config.xlsxFileName)
    elens_list, le_list = info_from_xlsx(xlsfile)
    del xlsfile
    aperture_list = Config.realApertureDiameter

    workpath = os.path.join(os.path.curdir, Config.workPath) 
    os.chdir(workpath)

    print_info(u"\n说明：本程序用来提取main.py的计算结果")
    print_info(u"\n      最终生成一系列的excel文件\n\n")

    print_info(u"按回车键来检查数据文件...", ",")
    raw_input()
    data = readResultFiles(elens_list, le_list, aperture_list)
    print_info(u"工作距离参数 %s" % "OK" if len(data)==len(elens_list) else "Error")
    print_info(u"落点能量参数 %s" % "OK" if len(data[0])==len(le_list) else "Error")
    print_info(u"光阑大小参数 %s" % "OK" if len(data[0][0])==len(aperture_list) else "Error")

    print_info(u"\n请选择，表格的第一行采用哪一个变量？")
    for i in range(3):
        print_info("    %s %s" % (i, vals[i]))
    first = -1
    while first < 0 or first >2:
        print u"请输入数字并回车：",
        try:
            first = int(raw_input())
        except:
            first = -1

    print_info(u"\n请选择，表格的第一列采用哪一个变量？")
    for i in range(3):
        if i != first:
            print_info("    %s %s" % (i, vals[i]))
    second = -1
    while second < 0 or second >2 or second == first:
        print u"请输入数字并回车：",
        try:
            second = int(raw_input())
        except:
            second = -1

    print first, second
    if first + second == 1:
        print_info(u"\n将生成%s个 .txt文件: 光阑大小****.txt\n" % len(aperture_list))
        print_info(u"按回车键继续...", ",")
        raw_input()
        dealwith_aperture(first, second, elens_list, le_list, aperture_list, data)
    elif first+second == 2:
        print_info(u"\n将生成%s个 .txt 文件: 落点能量****.txt\n" % len(le_list))
        print_info(u"按回车键继续...", ",")
        raw_input()
        dealwith_le(first, second, elens_list, le_list, aperture_list, data)
    else:
        print_info(u"\n将生成%s个 .txt 文件: 工作距离****.txt\n" % len(elens_list))
        print_info(u"按回车键继续...", ",")
        raw_input()
        dealwith_wd(first, second, elens_list, le_list, aperture_list, data)

    print_info(u"\n提取结束")
    raw_input()
    print u"是否提取激励和灵敏度(y/n)：",
    ans = raw_input()
    if ans.startswith('y'):
        print_info(u"\n请选择，表格的第一行采用哪一个变量？")
        for i in range(2):
            print_info("    %s %s" % (i, vals[i]))
        first = -1
        while first < 0 or first >1:
            print u"请输入数字并回车：",
            try:
                first = int(raw_input())
            except:
                first = -1
        extract_other_two(first, elens_list, le_list, data)
        print u"\n提取结束"
    raw_input()

